package cn.jbolt.ai.core.chain;

import cn.hutool.core.util.StrUtil;
import cn.jbolt.ai.app.entity.AiApp;
import cn.jbolt.ai.app.entity.AiAppMessage;
import cn.jbolt.ai.app.entity.AiAppSession;
import cn.jbolt.ai.app.vo.AiChatData;
import cn.jbolt.ai.app.vo.ActionData;
import cn.jbolt.ai.core.chain.node.base.RenderType;
import cn.jbolt.ai.core.chain.source.SourceChannel;
import cn.jbolt.ai.core.chain.source.SourceType;
import cn.jbolt.ai.klb.entity.AiDoc;
import cn.jbolt.util.DateUtil;
import cn.jbolt.util.JSONUtil;
import com.fasterxml.jackson.databind.node.ArrayNode;

import java.util.HashMap;
import java.util.List;

public class ChainContext {

    private SourceType sourceType;


    /**
     * 请求消息
     */
    private AiAppMessage requestMsg;
    /**
     * 响应消息
     */
    private AiAppMessage responseMsg;

    private AiApp aiApp;

    private AiAppSession aiSession;

    private ActionData actionData;


    private HashMap<String, Object> args = new HashMap<>();

    private SourceChannel sourceChannel;


    /**
     * 发送结束消息
     */
    public void sendOverMsg() {
        ActionData msg = ActionData.copy(actionData);
        msg.setAction(ActionData.ACTION_OVER);
        if (sourceChannel != null) {
            sourceChannel.sendSocketData(msg);
        }
    }

    /**
     * 发送同步数据消息
     */
    public void sendSyncDataMsg() {
        ActionData msg = ActionData.copy(actionData);
        msg.setAction(ActionData.ACTION_SYNCDATA);
        msg.setData(JSONUtil.createObjectNode(
                "request", JSONUtil.createObjectNode("id", requestMsg.getId(), "createTime", DateUtil.format(requestMsg.getCreateTime())),
                "response", JSONUtil.createObjectNode("id", responseMsg.getId(), "createTime", DateUtil.format(responseMsg.getCreateTime()))));
        if (sourceChannel != null) {
            sourceChannel.sendSocketData(msg);
        }
    }

    /**
     * 发送响应消息
     */
    public void sendResponseMsg(String msg, RenderType renderType) {
        ActionData data = ActionData.copy(actionData);
        data.setAction(ActionData.ACTION_RESPONSE);
        AiChatData assistantMsg = AiChatData.createAssistantMsg(msg, renderType);
        data.setData(assistantMsg);
        if (sourceChannel != null) {
            sourceChannel.sendSocketData(data);
        }
    }

    /**
     * 发送异常消息
     */
    public void sendErrorMsg(String errorMsg) {
        ActionData msg = ActionData.copy(actionData);
        msg.setAction(ActionData.ACTION_ERROR);
        msg.setData(new AiChatData().setMsg(errorMsg));
        if (sourceChannel != null) {
            sourceChannel.sendSocketData(msg);
        }
    }

    /**
     * 发送参考消息
     *
     * @param list
     */
    public void sendReference(List<AiDoc> list) {
        if (list == null || list.isEmpty()) return;
        ArrayNode array = JSONUtil.createArrayNode();
        for (AiDoc aiDoc : list) {
            array.add(aiDoc.convertToReference());
        }
        ActionData data = ActionData.copy(actionData);
        data.setAction(ActionData.ACTION_REFERENCE);
        data.setData(array);
        if (sourceChannel != null) {
            sourceChannel.sendSocketData(data);
        }
    }


    public void sendSuccessMsgTip(String content) {
        sendResponseMsg(StrUtil.format("<msg-tip type='success'>{}</msg-tip>\n\n", content), RenderType.markdown);
    }

    public void sendInfoMsgTip(String content) {
        sendResponseMsg(StrUtil.format("<msg-tip type='info'>{}</msg-tip>\n\n", content), RenderType.markdown);
    }

    public void sendErrorMsgTip(String content) {
        sendResponseMsg(StrUtil.format("<msg-tip type='error'>{}</msg-tip>\n\n", content), RenderType.markdown);
    }


    public ActionData getActionData() {
        return actionData;
    }

    public void setActionData(ActionData actionData) {
        this.actionData = actionData;
    }

    public AiApp getAiApp() {
        return aiApp;
    }

    public void setAiApp(AiApp aiApp) {
        this.aiApp = aiApp;
    }

    public AiAppSession getAiSession() {
        return aiSession;
    }

    public void setAiSession(AiAppSession aiSession) {
        this.aiSession = aiSession;
    }


    public <T> T getAttr(String key) {
        if (StrUtil.isBlank(key)) {
            return null;
        }
        return (T) this.args.get(key);
    }

    public ChainContext setAttr(String key, Object value) {
        this.args.put(key, value);
        return this;
    }

    public HashMap<String, Object> getArgs() {
        return args;
    }

    public AiAppMessage getRequestMsg() {
        return requestMsg;
    }

    public void setRequestMsg(AiAppMessage requestMsg) {
        this.requestMsg = requestMsg;
    }

    public AiAppMessage getResponseMsg() {
        return responseMsg;
    }

    public void setResponseMsg(AiAppMessage responseMsg) {
        this.responseMsg = responseMsg;
    }

    public SourceType getSourceType() {
        return sourceType;
    }

    public void setSourceType(SourceType sourceType) {
        this.sourceType = sourceType;
    }

    public SourceChannel getSourceChannel() {
        return sourceChannel;
    }

    public void setSourceChannel(SourceChannel sourceChannel) {
        this.sourceChannel = sourceChannel;
    }
}
